#pragma once

#include "rewine/MapIndicator.hpp"
#include "rewine/impl/CImage.hpp"
#include "rewine/impl/CPeExport.hpp"
#include "rewine/impl/CPeImport.hpp"

#include "wininc/pe.h"

namespace rewine
{

class ElfInterop;

class CPeImage : public CImage {
private:

    std::shared_ptr<MapIndicator>   m_spSourceMap;
    bool                            m_bNeedRemap;
    std::shared_ptr<MapIndicator>   m_spMap;
    std::shared_ptr<MapIndicator>   m_spHeadersMap;

    PIMAGE_NT_HEADERS           nt_hdr;
    PIMAGE_FILE_HEADER          file_hdr;
    PIMAGE_OPTIONAL_HEADER      opt_hdr;

    int                                         nbSection;
    PIMAGE_SECTION_HEADER                       sectionHeaders;
    std::vector<std::shared_ptr<MapIndicator>>  sectionMaps;

    int                                         nbExport;
    ULONG                                       szExportd;
    PIMAGE_EXPORT_DIRECTORY                     exportd;
    std::vector<std::shared_ptr<CPeExport>>     exports;

    int                                         nbImport;
    std::vector<std::shared_ptr<CPeImportLib>>  imports;

    int                                         nbDelayload;
    std::vector<std::shared_ptr<CPeImportLib>>  delayloads;

public:

    CPeImage() = default;
    ~CPeImage() override = default;

public:

    PVOID rva2va(LONG_PTR offset, const MapIndicator &map) const;
    PVOID rva2va(LONG_PTR offset) const;

    PVOID rva2fp(PIMAGE_SECTION_HEADER secHdr, LONG_PTR offset, const MapIndicator &map) const;
    PVOID rva2fp(PIMAGE_SECTION_HEADER secHdr, LONG_PTR offset) const;

    DWORD directory_data_rva(USHORT index, OUT PULONG size) const;
    PVOID directory_data(USHORT index, OUT PULONG size) const;
    PIMAGE_SECTION_HEADER find_section_header_by_rva(DWORD vma) const;

private:

    int parse_dos_header_and_goto_nt(MapIndicator &map, size_t offset = 0);
    int use_nt_headers_and_next(MapIndicator &map, size_t offset = (size_t)-1);
    int check_headers_compatibility();
    int use_section_headers_and_next(MapIndicator &map, size_t offset = (size_t)-1);
    int parse_dll_name();
    bool is_forwarder_export(LONG_PTR rva) const;

    int map_base();
    int map_headers();
    int map_sections();
    int relocate();

    int load_exports();

    int load_imports();
    int load_import(int index, PIMAGE_IMPORT_DESCRIPTOR impd, CPeImportLib& dll);
    int load_delayload_imports();
    int load_delayload_import(int index, PIMAGE_DELAYLOAD_DESCRIPTOR delayd, CPeImportLib& dll);

    int bind_import_symbol(CPeImportLib& dll, int index, PIMAGE_THUNK_DATA pName, PIMAGE_THUNK_DATA pAddress,
                           CPeImport& imp, bool bBindNow);

    int bind_forwarder_exports();
    int bind_forwarder_export(CPeExport& exp);

    int set_memory_acl();

public:

    int initFromMap(const std::shared_ptr<MapIndicator> &map, bool bRemap);
    int initFromFd(int fd);

    int parse();
    int load();
    int unload();

public:

    ImageFlavour flavour() const override {
        return ImageFlavour::PE;
    }

    int getExportCount() const override;
    IExport * getExportByIndex(unsigned int index) const override;
    IExport * getExportByOrdinal(DWORD ordinal) const;
    IExport * getExportByName(const std::string& name, uint16_t hint = -1) const override;

    PVOID getImageBase() const {
        return (PVOID) m_spMap->start;
    }

    size_t getImageSize() const {
        return opt_hdr->SizeOfImage;
    }

    PVOID getEntryPoint() const {
        return rva2va(opt_hdr->AddressOfEntryPoint);
    }

    using dllmain_t = BOOL(*)(PVOID, DWORD, PVOID);
    BOOL onEvent(DWORD dwReason) const override;

    friend ::rewine::ElfInterop;

};

} // namespace rewine
